from loadDataset import LoadBCICompDataSet
from modelSource import Multi_featureV2_modify as multi_person_feature
from loadDataset import LoadModelParam
import torch
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

modelPath = r""
filePath = r""

def predict():

    model = LoadModelParam.LoadModelParam(modelPath, multi_person_feature.multi_person_feature(1, 2)).GetModel().to(device)

    datas = LoadBCICompDataSet.LoadBCICompDataSet(filePath)

    ans = []

    for epoch in range(0, np.size(datas['Signal'], 0)):
        m = 0
        data = torch.from_numpy(datas['Responses'][epoch]).float().to(device)
        for i in range(0, data.size()[1]):
            for j in range(0, data.size()[0]):
                _input = data[j, i, :, 61:64].unsqueeze(0).unsqueeze(0)
                _input = torch.cat((_input, _input))
                _input = torch.cat((_input, _input), dim=3)
                output = model(_input)
                if output.argmax() == 1:
                    ans.append(j)

    print(ans)
